import os
import gradio as gr
from configs.model_config import *
from chains.chatbot import ChatBot


# LLM 使用的历史对话长度
LLM_HISTORY_LEN = 3


# 模型列表
prompt_model_dict_list = list(prompt_model_dict.keys())
llm_model_dict_list = list(llm_model_dict.keys())
# 创建机器人
chat_bot = ChatBot()
# 初始化模型
def init_model():
    try:
        chat_bot.init_cfg()
        return """模型已成功加载"""
    except Exception as e:
        print(e)
        return """模型加载失败，请重新选择"""


# 重载模型
def reinit_model(prompt_model, llm_history_len, llm_model, history):
    try:
        chat_bot.init_cfg(
            prompt_model=prompt_model,
            llm_history_len=llm_history_len,
            llm_model=llm_model
        )
        model_status = """模型重新加载成功"""
    except Exception as e:
        print(e)
        model_status = """模型未成功重载，请重新选择"""
    return history + [[None, model_status]]


webui_title = """
# Chat-AIGC WebUI
"""

# css 样式
block_css = """.importantButton {
    background: linear-gradient(45deg, #7e0570,#5d1c99, #6e00ff) !important;
    border: none !important;
}
.importantButton:hover {
    background: linear-gradient(45deg, #ff00e0,#8500ff, #6e00ff) !important;
    border: none !important;
}"""


# 加载模型
model_status = init_model()


with gr.Blocks(css=block_css) as demo:
    model_status = gr.State(model_status)
    gr.Markdown(webui_title)
    with gr.Row():
        with gr.Column(scale=2):
            # 清空了用户输入，并更新了 chatbot 的聊天记录
            def user(user_message, history):
                return "", history + [[user_message, None]]
            
            # bot 基于历史记录，产生最新的答复
            def bot(history):
                history = chat_bot.chat(history)
                return history

            chatbot = gr.Chatbot(
                [[None, model_status.value]],
                elem_id="chat-box",
                show_label=False
            )
            query = gr.Textbox(
                show_label=False,
                placeholder="'/image + prompt' can use sd-model to create picture for you."
            ).style(container=False)
            query.submit(user, [query, chatbot], [query, chatbot], queue=False).then(
                bot, chatbot, chatbot
            )
        with gr.Column(scale=1):
            llm_model = gr.Radio(
                llm_model_dict_list,
                label="LLM 模型",
                value=LLM_MODEL,
                interactive=True
            )
            llm_history_len = gr.Slider(
                1, 5,
                value=LLM_HISTORY_LEN,
                step=1,
                label="LLM history len",
                interactive=True
            )
            prompt_model = gr.Radio(
                prompt_model_dict_list,
                label="Prompt生成模型",
                value=PROMPT_MODEL,
                interactive=True
            )
            load_model_button = gr.Button("重新加载模型")
    load_model_button.click(
        reinit_model,
        show_progress=True,
        inputs=[prompt_model, llm_history_len, llm_model, chatbot],
        outputs=chatbot
    )


demo.launch(
    debug=True,
    enable_queue=True,
    server_name="0.0.0.0",
    server_port=7862,
    share=False,
    inbrowser=False
)
